You write custom CUDA kernels to replace the PyTorch operators in the given EvoNorm architecture to get speedups.
You have complete freedom to choose the set of operators you want to replace. You may make the decision to replace some operators with custom CUDA kernels and leave others unchanged. You may replace multiple operators with custom implementations, consider operator fusion opportunities (combining multiple operators into a single kernel, for example, combining normalization+affine_transform+nonlinear_gating), or algorithmic changes (such as optimized memory access patterns). You are only limited by your imagination.

Technologies Used :

PyTorch: Deep learning framework

CUDA: GPU acceleration for parallel computing

C++/CUDA C++: High-performance kernel programming

Inline C++/CUDA Extension (torch.utils.cpp_extension.load_inline): Just-In-Time (JIT) compilation of custom operators

Minkowski Distance: Generalized distance metric with parameter p

Vectorized Memory Access (float4): Uses 128-bit wide loads (4 floats) to improve memory bandwidth utilization

Two-Dimensional Grid Layout: Uses dim3(N, blocks_per_instance) for parallel processing across batches and feature dimensions

Warp-Level Primitives: Uses __shfl_down_sync for efficient intra-warp reduction

Two-Stage Parallel Reduction: Combines warp-level reduction with shared memory and block-level reduction

Strided Memory Access: Threads process elements with calculated stride for load distribution

Atomic Operations (atomicAdd): Safely accumulates results from multiple thread blocks

Mathematical Operations: Uses powf, fabsf for Minkowski distance calculation

Fast Math Operations: Uses --use_fast_math compiler flag for optimized mathematical functions

Configurable Parallelism: blocks_per_instance parameter controls the degree of parallelism per sample

Memory Coalescing: Optimized memory access patterns through contiguous tensor layout and vectorized loads

Numerical Stability: Adds epsilon (eps) to prevent numerical issues in power operations

Flexible Distance Metric: Supports arbitrary p-values for generalized Minkowski distance

Tensor Contiguity Enforcement: Ensures optimal memory layout in PyTorch wrapper

Automatic Device Placement: Ensures tensors are on CUDA device

Efficient Reduction Pattern: Implements hierarchical reduction from thread to warp to block level


Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is:
import torch
import torch.nn as nn
import torch.nn.functional as F

N, C, H, W = 32, 64, 56, 56
EPS = 1e-6


class MinkowskiDistance(nn.Module):

    def __init__(self, p=2.0, keepdim=False, eps=1e-6):
        super().__init__()
        self.p = p
        self.keepdim = keepdim
        self.eps = eps

    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        diff = torch.abs(x - y)

        pow_diff = torch.pow(diff, self.p)

        sum_pow = torch.sum(pow_diff, dim=[1, 2, 3], keepdim=self.keepdim)

        output = torch.pow(sum_pow + self.eps, 1.0 / self.p)

        return output


class Model(nn.Module):

    def __init__(self, p=3.0):
        super().__init__()
        self.op = MinkowskiDistance(p=p, keepdim=False, eps=EPS)

    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        return self.op(x, y)


def get_inputs():
    x = torch.randn(N, C, H, W, dtype=torch.float32)
    y = torch.randn(N, C, H, W, dtype=torch.float32)
    return [x, y]


def get_init_inputs():
    p = 3.0
    return [p]